# tcmsp_extractor.py
import pandas as pd
import requests
import re
import os
import urllib.parse
import json
import time

def get_current_token():
    """
    Extract the current token from the TCMSP search page
    
    Returns:
        str or None: The current token if found, None otherwise
    """
    try:
        search_url = "https://www.tcmsp-e.com/tcmspsearch.php"
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
        }
        
        response = requests.get(search_url, headers=headers, timeout=30)
        
        if response.status_code == 200:
            # Look for the token in the form
            token_match = re.search(r'<input type="hidden" name="token" value="([a-f0-9]{32})">', response.text)
            if token_match:
                token = token_match.group(1)
                print(f"Found current token: {token}")
                return token
        
        print("Failed to extract token from search page")
        return None
    except Exception as e:
        print(f"Error getting token: {e}")
        return None

def extract_tcmsp_data_by_en_name(herb_en_name):
    """
    Access TCMSP website and extract data using English name
    
    Args:
        herb_en_name (str): English name of the herb to search for
        
    Returns:
        pandas.DataFrame or None: DataFrame containing the herb data, or None if extraction failed
    """
    # Get current token
    token = get_current_token()
    if not token:
        print("WARNING: Could not get valid token, aborting request")
        return None
    
    # Encode the herb name for URL
    encoded_herb = urllib.parse.quote(herb_en_name)
    
    # Construct the URL with current token
    url = f"https://www.tcmsp-e.com/tcmspsearch.php?qr={encoded_herb}&qsr=herb_en_name&token={token}"
    
    print(f"Accessing URL with English name: {url}")
    
    # Send the request with headers to mimic a browser
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
    }
    
    try:
        response = requests.get(url, headers=headers, timeout=30)
        
        if response.status_code != 200:
            print(f"Failed to access URL for {herb_en_name}. Status code: {response.status_code}")
            return None
        
        # Extract Chinese name of the herb for verification
        chinese_name_match = re.search(r'"herb_cn_name":"([^"]+)"', response.text)
        if chinese_name_match:
            chinese_name = chinese_name_match.group(1)
            pinyin_match = re.search(r'"herb_pinyin":"([^"]+)"', response.text)
            pinyin = pinyin_match.group(1) if pinyin_match else ""
            print(f"Found herb: {chinese_name} ({pinyin})")
        else:
            print(f"Could not find Chinese name for {herb_en_name}")
        
        # Extract the data part directly from source HTML
        pattern = r'\$\("#grid"\)\.kendoGrid\(\{\s+dataSource:\s+\{\s+data:\s+(\[.*?\]),\s+pageSize:'
        data_match = re.search(pattern, response.text, re.DOTALL)
        
        if not data_match:
            print(f"Could not find grid data pattern for {herb_en_name}")
            return None
        
        data_json = data_match.group(1)
        
        try:
            data = json.loads(data_json)
            print(f"Successfully extracted data with {len(data)} rows")
            
            # Convert to DataFrame
            df = pd.DataFrame(data)
            
            # Modify column names and drop specific columns
            if not df.empty:
                
                # Rename columns
                df = df.rename(columns={
                    "molecule_ID": "化合物ID",
                    "molecule_name": "化合物名称"
                })
                
                # Drop 2nd, 3rd, and 7th columns (indices 1, 2, 6)
                if len(df.columns) > 6:  # Ensure we have enough columns
                    df = df.drop(df.columns[[1, 2, 6]], axis=1)
            
            return df
            
        except json.JSONDecodeError as e:
            print(f"Error parsing JSON for {herb_en_name}: {e}")
            print(f"JSON snippet: {data_json[:100]}...")
            return None
            
    except Exception as e:
        print(f"Error processing {herb_en_name}: {e}")
        return None

def extract_tcmsp_data_by_cn_name(herb_cn_name):
    """
    Access TCMSP website and extract data using Chinese name
    
    Args:
        herb_cn_name (str): Chinese name of the herb to search for
        
    Returns:
        tuple: (pandas.DataFrame or None, str or None) - DataFrame containing the herb data and English name if found
    """
    # Get current token
    token = get_current_token()
    if not token:
        print("WARNING: Could not get valid token, aborting request")
        return None, None
    
    # Encode the herb name for URL
    encoded_herb = urllib.parse.quote(herb_cn_name)
    
    # Construct the URL for searching by herb name (any language) with current token
    url = f"https://www.tcmsp-e.com/tcmspsearch.php?qs=herb_all_name&q={encoded_herb}&token={token}"
    
    print(f"Accessing URL with Chinese name: {url}")
    
    # Send the request with headers to mimic a browser
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
    }
    
    try:
        response = requests.get(url, headers=headers, timeout=30)
        
        if response.status_code != 200:
            print(f"Failed to access URL for {herb_cn_name}. Status code: {response.status_code}")
            return None, None
        
        # Check if we got a table of herb names instead of ingredients
        herb_grid_match = re.search(r'\$\("#grid"\)\.kendoGrid\(\{\s+dataSource:\s+\{\s+data:\s+(\[{"herb_cn_name".*?\]),\s+pageSize:', response.text, re.DOTALL)
        
        if herb_grid_match:
            # We have a list of herbs matching the search
            herbs_json = herb_grid_match.group(1)
            try:
                herbs_data = json.loads(herbs_json)
                print(f"Found {len(herbs_data)} herbs matching '{herb_cn_name}'")
                
                # Get the English name of the first herb
                if len(herbs_data) > 0 and "herb_en_name" in herbs_data[0]:
                    eng_name = herbs_data[0]["herb_en_name"]
                    print(f"Using English name: {eng_name}")
                    
                    # Wait a bit before the next request
                    time.sleep(3)
                    
                    # Now search by English name
                    return extract_tcmsp_data_by_en_name(eng_name), eng_name
                else:
                    print(f"No English name found for {herb_cn_name}")
                    return None, None
                    
            except json.JSONDecodeError as e:
                print(f"Error parsing herbs JSON for {herb_cn_name}: {e}")
                return None, None
        
        # If we get here, try to extract ingredient data directly
        ingredient_pattern = r'\$\("#grid"\)\.kendoGrid\(\{\s+dataSource:\s+\{\s+data:\s+(\[.*?\]),\s+pageSize:'
        data_match = re.search(ingredient_pattern, response.text, re.DOTALL)
        
        if not data_match:
            print(f"Could not find grid data pattern for {herb_cn_name}")
            return None, None
        
        data_json = data_match.group(1)
        
        try:
            data = json.loads(data_json)
            print(f"Successfully extracted data with {len(data)} rows")
            
            # Convert to DataFrame
            df = pd.DataFrame(data)
            
            # Modify column names and drop specific columns
            if not df.empty:
                # Print original columns for verification
                print(f"Original column names: {df.columns.tolist()}")
                
                # Rename columns
                # df = df.rename(columns={
                #     "molecule_ID": "化合物ID",
                #     "molecule_name": "化合物名称"
                # })
                
                # Drop 2nd, 3rd, and 7th columns (indices 1, 2, 6)
                # if len(df.columns) > 6:  # Ensure we have enough columns
                #     df = df.drop(df.columns[[1, 2, 6]], axis=1)
            
            # Try to extract English name from the page
            chinese_name_match = re.search(r'"herb_cn_name":"([^"]+)"', response.text)
            eng_name_match = re.search(r'"herb_en_name":"([^"]+)"', response.text)
            
            eng_name = None
            if eng_name_match:
                eng_name = eng_name_match.group(1)
                
            return df, eng_name
            
        except json.JSONDecodeError as e:
            print(f"Error parsing JSON for {herb_cn_name}: {e}")
            print(f"JSON snippet: {data_json[:100]}...")
            return None, None
            
    except Exception as e:
        print(f"Error processing {herb_cn_name}: {e}")
        return None, None

def process_herb_list_from_excel(excel_file, output_dir="tcmsp_results"):
    """
    Process a list of herbs from an Excel file
    
    Args:
        excel_file (str): Path to Excel file containing herb names
        output_dir (str, optional): Directory to save results
        
    Returns:
        tuple: (successful_count, failed_count)
    """
    # Create output directory if it doesn't exist
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # Read the Excel file
    try:
        herb_df = pd.read_excel(excel_file)
        print(f"Successfully read Excel file with {len(herb_df)} rows")
    except Exception as e:
        print(f"Error reading Excel file: {e}")
        return (0, 0)
    
    # Check if required columns exist
    if "中文名" not in herb_df.columns:
        print("Required column '中文名' not found in the Excel file.")
        print(f"Available columns: {herb_df.columns.tolist()}")
        return (0, 0)
    
    # Create a copy of the DataFrame to update English names
    updated_df = herb_df.copy()
    
    # Process each row
    successful = 0
    failed = 0
    
    # Inside process_herb_list_from_excel function
    for index, row in herb_df.iterrows():
        chinese_name = row['中文名']
        english_name = row['英文名'] if '英文名' in herb_df.columns and not pd.isna(row['英文名']) else None
    
        print(f"\nProcessing {index+1}/{len(herb_df)}: {chinese_name} (English: {english_name if english_name else 'Not available'})")
    
        result_df = None
        eng_name_found = None
    
        # First try with English name if available
        if english_name and english_name.strip() != "":
            result_df = extract_tcmsp_data_by_en_name(english_name)
    
        # If no result or English name not available, try with Chinese name
        if result_df is None or result_df.empty:
            print(f"Trying with Chinese name: {chinese_name}")
            result_df, eng_name_found = extract_tcmsp_data_by_cn_name(chinese_name)
        
            # Update the English name in our DataFrame if we found one through Chinese name search
            # This modification will update any herb that required Chinese name search
            if eng_name_found:
                print(f"Updating English name for {chinese_name} to: {eng_name_found}")
                updated_df.loc[index, '英文名'] = eng_name_found
        
        # Save to Excel if data is found
        if result_df is not None and not result_df.empty:
            output_file = os.path.join(output_dir, f"{chinese_name}.xlsx")
            result_df.to_excel(output_file, index=False, engine='openpyxl')
            print(f"  Data saved to {output_file} with {len(result_df)} rows")
            successful += 1
        else:
            print(f"  No data found for {chinese_name}")
            failed += 1
        
        # Add a delay between herbs to avoid overloading the server
        if index < len(herb_df) - 1:  # Don't wait after the last item
            wait_time = 5  # 5 seconds between herbs
            print(f"Waiting {wait_time} seconds before processing next herb...")
            time.sleep(wait_time)
    
    # Save the updated DataFrame with new English names
    if '英文名' in updated_df.columns and updated_df['英文名'].notna().sum() > herb_df['英文名'].notna().sum():
        updated_excel = os.path.join(output_dir, "updated_" + os.path.basename(excel_file))
        updated_df.to_excel(updated_excel, index=False, engine='openpyxl')
        print(f"\nUpdated Excel file with new English names saved to: {updated_excel}")
    
    print(f"\nProcessing complete. Summary: {successful} successful, {failed} failed.")
    return (successful, failed)

# Only run this if the script is executed directly
if __name__ == "__main__":
    process_herb_list_from_excel("./API/中药.xlsx")